/*
 * SPDX-FileCopyrightText: 2025 Espressif Systems (Shanghai) CO LTD
 *
 * SPDX-License-Identifier: Apache-2.0
 */

#include <errno.h>
#include <vector>
#include <cstring>
#include "freertos/FreeRTOS.h"
#include "freertos/task.h"
#include "esp_wifi.h"
#include "esp_event.h"
#include "esp_netif.h"
#include "esp_log.h"
#include "esp_mac.h"
#include "cJSON.h"
#include "lwip/sockets.h"
#include "lwip/ip_addr.h"
#include "lwip/inet.h"
#include "app_ap_conf.hpp"

namespace esp_brookesia::speaker_apps {

static const char *TAG = "ApProvision";

/* Symbols generated by EMBED_TXTFILES. */
extern const char wifi_html_start[] asm("_binary_wifi_html_start");
extern const char wifi_html_end[]   asm("_binary_wifi_html_end");

/* Forward declarations */
static void captive_dns_task(void *arg);

/*********************** Public API ****************************/

esp_err_t ApProvision::start(
    const CredentialsCallback &cb,
    const StateChangeCallback &sc_cb,
    const std::vector<wifi_ap_record_t> &initial_aps
)
{
    if (_running) {
        /* Already running, just update callback if provided */
        if (cb) {
            _cb = cb;
        }
        if (sc_cb) {
            _sc_cb = sc_cb;
        }
        return ESP_OK;
    }

    _cb = cb;
    _sc_cb = sc_cb;

    {
        std::lock_guard<std::mutex> lock(_initial_aps_mutex);
        _initial_aps = initial_aps;
    }

    if (_sc_cb) {
        _sc_cb(true);
    }


    ESP_ERROR_CHECK_WITHOUT_ABORT(esp_wifi_get_mode(&_previous_mode));

    /* Create AP netif if not already */
    if (_ap_netif == nullptr) {
        _ap_netif = esp_netif_create_default_wifi_ap();
        if (_ap_netif == nullptr) {
            ESP_LOGE(TAG, "Failed to create default AP netif");
            return ESP_FAIL;
        }
        /* Ensure DHCP server running */
        esp_netif_dhcps_start(_ap_netif);
    }

    /* Enable AP+STA */
    ESP_ERROR_CHECK_WITHOUT_ABORT(esp_wifi_set_mode(WIFI_MODE_APSTA));
    ESP_ERROR_CHECK(init_softap());

    /* Register event handlers */
    ESP_ERROR_CHECK_WITHOUT_ABORT(
        esp_event_handler_register(WIFI_EVENT, ESP_EVENT_ANY_ID, &ApProvision::wifi_event_handler, nullptr));
    ESP_ERROR_CHECK_WITHOUT_ABORT(
        esp_event_handler_register(IP_EVENT, IP_EVENT_STA_GOT_IP, &ApProvision::ip_event_handler, nullptr));

    /* Start captive portal services */
    ESP_ERROR_CHECK(init_dns_server());
    ESP_ERROR_CHECK(init_http_server());

    _running = true;
    ESP_LOGI(TAG, "Provisioning started, AP SSID: %s", _ap_ssid);
    return ESP_OK;
}

esp_err_t ApProvision::stop()
{
    if (!_running) {
        return ESP_OK;
    }

    ESP_LOGI(TAG, "Stopping provisioning");

    deinit_http_server();
    deinit_dns_server();

    esp_event_handler_unregister(WIFI_EVENT, ESP_EVENT_ANY_ID, &ApProvision::wifi_event_handler);
    esp_event_handler_unregister(IP_EVENT, IP_EVENT_STA_GOT_IP, &ApProvision::ip_event_handler);

    ESP_ERROR_CHECK_WITHOUT_ABORT(esp_wifi_set_mode(WIFI_MODE_STA));

    if (_sc_cb) {
        _sc_cb(false);
    }
    _running = false;
    _target_ssid.clear();
    _target_password.clear();

    return ESP_OK;
}

const char *ApProvision::get_ap_ssid()
{
    return _running ? _ap_ssid : nullptr;
}

void ApProvision::register_callback(const CredentialsCallback &cb)
{
    _cb = cb;
}

void ApProvision::update_ap_list(const std::vector<wifi_ap_record_t> &aps)
{
    std::lock_guard<std::mutex> lock(_initial_aps_mutex);
    _initial_aps = aps;
}

/*********************** Internal helpers ****************************/

esp_err_t ApProvision::init_softap()
{
    /* Generate SSID "ESP-Brookesia-XXXX" */
    uint8_t mac[6];
    esp_read_mac(mac, ESP_MAC_WIFI_SOFTAP);
    snprintf(_ap_ssid, sizeof(_ap_ssid), "ESP-Brookesia-%02X%02X", mac[4], mac[5]);

    wifi_config_t ap_cfg = {};
    strncpy((char *)ap_cfg.ap.ssid, _ap_ssid, sizeof(ap_cfg.ap.ssid));
    ap_cfg.ap.ssid_len = strlen(_ap_ssid);
    ap_cfg.ap.channel = 1;
    ap_cfg.ap.max_connection = 4;
    ap_cfg.ap.authmode = WIFI_AUTH_OPEN; // No password
    ap_cfg.ap.ssid_hidden = 0; // Visible when provisioning is active

    ESP_ERROR_CHECK_WITHOUT_ABORT(esp_wifi_set_config(WIFI_IF_AP, &ap_cfg));
    return ESP_OK;
}

/*********************** HTTP server ****************************/

static esp_err_t http_send_file(httpd_req_t *req, const char *start, const char *end, const char *content_type)
{
    httpd_resp_set_type(req, content_type);
    size_t len = end - start;
    return httpd_resp_send(req, start, len);
}

esp_err_t ApProvision::handle_root(httpd_req_t *req)
{
    return http_send_file(req, wifi_html_start, wifi_html_end, "text/html");
}

esp_err_t ApProvision::handle_scan(httpd_req_t *req)
{
    std::vector<wifi_ap_record_t> records;

    /* Perform a synchronous Wi-Fi scan and return fresh results */
    wifi_scan_config_t scan_cfg = {};
    scan_cfg.ssid = nullptr;
    scan_cfg.bssid = nullptr;
    scan_cfg.channel = 0;
    scan_cfg.show_hidden = false;
    ESP_LOGI(TAG, "Starting synchronous Wi-Fi scan");
    if (esp_wifi_scan_start(&scan_cfg, true) == ESP_OK) {
        uint16_t ap_num = 0;
        if (esp_wifi_scan_get_ap_num(&ap_num) == ESP_OK && ap_num > 0) {
            ESP_LOGI(TAG, "Scan done, %d APs found", ap_num);
            records.resize(ap_num);
            esp_wifi_scan_get_ap_records(&ap_num, records.data());
        }
    }
    /* Cache latest scan list for potential future use */
    {
        std::lock_guard<std::mutex> lock(_initial_aps_mutex);
        _initial_aps = records;
    }

    if (records.empty()) {
        httpd_resp_set_type(req, "application/json");
        httpd_resp_sendstr(req, "[]");
        return ESP_OK;
    }

    /* Build JSON */
    cJSON *root = cJSON_CreateArray();
    for (const auto &record : records) {
        ESP_LOGD(TAG, "AP: %s, %d dBm", record.ssid, record.rssi);
        cJSON *item = cJSON_CreateObject();
        cJSON_AddStringToObject(item, "ssid", (char *)record.ssid);
        cJSON_AddNumberToObject(item, "rssi", record.rssi);
        cJSON_AddItemToArray(root, item);
    }
    char *json_str = cJSON_PrintUnformatted(root);
    cJSON_Delete(root);

    httpd_resp_set_type(req, "application/json");
    esp_err_t ret = httpd_resp_sendstr(req, json_str);
    cJSON_free(json_str);
    return ret;
}

esp_err_t ApProvision::handle_connect(httpd_req_t *req)
{
    /* Read body */
    char buf[256];
    int recv = httpd_req_recv(req, buf, sizeof(buf));
    if (recv <= 0) {
        return httpd_resp_send_err(req, HTTPD_400_BAD_REQUEST, "Invalid body");
    }
    buf[recv] = '\0';

    cJSON *root = cJSON_Parse(buf);
    if (!root) {
        return httpd_resp_send_err(req, HTTPD_400_BAD_REQUEST, "JSON parse error");
    }
    const cJSON *ssid_item = cJSON_GetObjectItem(root, "ssid");
    const cJSON *pwd_item  = cJSON_GetObjectItem(root, "password");

    if (!cJSON_IsString(ssid_item) || ssid_item->valuestring == nullptr) {
        cJSON_Delete(root);
        return httpd_resp_send_err(req, HTTPD_400_BAD_REQUEST, "SSID missing");
    }

    std::string ssid = ssid_item->valuestring;
    std::string pwd  = (pwd_item && cJSON_IsString(pwd_item)) ? pwd_item->valuestring : "";
    cJSON_Delete(root);

    ESP_LOGI(TAG, "User choose SSID: %s", ssid.c_str());

    /* Keep target for later verification */
    _target_ssid = ssid;
    _target_password = pwd;

    /* Update connection status */
    _connect_status = ConnectStatus::CONNECTING;
    _connect_error_msg.clear();

    wifi_config_t sta_cfg = {};
    strncpy((char *)sta_cfg.sta.ssid, ssid.c_str(), sizeof(sta_cfg.sta.ssid));
    strncpy((char *)sta_cfg.sta.password, pwd.c_str(), sizeof(sta_cfg.sta.password));

    esp_err_t err = esp_wifi_set_config(WIFI_IF_STA, &sta_cfg);
    if (err != ESP_OK) {
        return httpd_resp_send_err(req, HTTPD_500_INTERNAL_SERVER_ERROR, "Set config failed");
    }
    esp_wifi_disconnect();
    esp_wifi_connect();

    httpd_resp_set_type(req, "application/json");
    httpd_resp_sendstr(req, "{\"success\":true}");
    return ESP_OK;
}

esp_err_t ApProvision::handle_status(httpd_req_t *req)
{
    httpd_resp_set_type(req, "application/json");

    cJSON *root = cJSON_CreateObject();
    const char *status_str = "idle";
    switch (_connect_status) {
    case ConnectStatus::CONNECTING: status_str = "connecting"; break;
    case ConnectStatus::SUCCESS:    status_str = "success";    break;
    case ConnectStatus::FAILED:     status_str = "failed";     break;
    case ConnectStatus::IDLE:
    default:                       status_str = "idle";       break;
    }
    cJSON_AddStringToObject(root, "status", status_str);
    if (_connect_status == ConnectStatus::FAILED && !_connect_error_msg.empty()) {
        cJSON_AddStringToObject(root, "reason", _connect_error_msg.c_str());
    }
    char *json_str = cJSON_PrintUnformatted(root);
    cJSON_Delete(root);
    httpd_resp_sendstr(req, json_str);
    cJSON_free(json_str);
    return ESP_OK;
}

static const httpd_uri_t ROOT_GET = {
    .uri      = "/",
    .method   = HTTP_GET,
    .handler  = [](httpd_req_t *req){ return ApProvision::handle_root(req);},
    .user_ctx = nullptr
};

static const httpd_uri_t SCAN_GET = {
    .uri      = "/api/wifi/scan",
    .method   = HTTP_GET,
    .handler  = [](httpd_req_t *req){ return ApProvision::handle_scan(req);},
    .user_ctx = nullptr
};

static const httpd_uri_t CONNECT_POST = {
    .uri      = "/api/wifi/connect",
    .method   = HTTP_POST,
    .handler  = [](httpd_req_t *req){ return ApProvision::handle_connect(req);},
    .user_ctx = nullptr
};

static const httpd_uri_t STATUS_GET = {
    .uri      = "/api/wifi/status",
    .method   = HTTP_GET,
    .handler  = [](httpd_req_t *req){ return ApProvision::handle_status(req);},
    .user_ctx = nullptr
};

esp_err_t ApProvision::init_http_server()
{
    httpd_config_t cfg = HTTPD_DEFAULT_CONFIG();
    cfg.uri_match_fn = httpd_uri_match_wildcard;
    cfg.max_uri_handlers = 8;
    cfg.lru_purge_enable = true;
    if (httpd_start(&_httpd, &cfg) != ESP_OK) {
        ESP_LOGE(TAG, "Failed to start httpd");
        return ESP_FAIL;
    }

    httpd_register_uri_handler(_httpd, &ROOT_GET);
    httpd_register_uri_handler(_httpd, &SCAN_GET);
    httpd_register_uri_handler(_httpd, &CONNECT_POST);
    httpd_register_uri_handler(_httpd, &STATUS_GET);

    /* Wildcard redirect for captive portal (any other URI) */
    static const httpd_uri_t ANY_GET = {
        .uri      = "/*",
        .method   = HTTP_GET,
        .handler  = [](httpd_req_t *req)
        {
            httpd_resp_set_status(req, "302 Found");
            httpd_resp_set_hdr(req, "Location", "http://192.168.4.1/");
            httpd_resp_send(req, nullptr, 0);
            return ESP_OK;
        },
        .user_ctx = nullptr
    };
    httpd_register_uri_handler(_httpd, &ANY_GET);

    return ESP_OK;
}

void ApProvision::deinit_http_server()
{
    if (_httpd) {
        httpd_stop(_httpd);
        _httpd = nullptr;
    }
}

/*********************** DNS server (captive portal) ****************************/

// Extremely small DNS server implementation: every query returns 192.168.4.1
#define DNS_PORT 53

/* DNS header structure */
struct dns_header_t {
    uint16_t id;
    uint16_t flags;
    uint16_t qdcount;
    uint16_t ancount;
    uint16_t nscount;
    uint16_t arcount;
} __attribute__((packed));

static void captive_dns_task(void *arg)
{
    (void)arg;
    int sock = socket(AF_INET, SOCK_DGRAM, IPPROTO_UDP);
    if (sock < 0) {
        ESP_LOGE(TAG, "DNS socket create failed: %d", errno);
        vTaskDelete(nullptr);
        return;
    }
    ApProvision::_dns_sock = sock;

    struct sockaddr_in addr = {};
    addr.sin_family = AF_INET;
    addr.sin_port   = htons(DNS_PORT);
    addr.sin_addr.s_addr = htonl(INADDR_ANY);

    if (bind(sock, (struct sockaddr *)&addr, sizeof(addr)) < 0) {
        ESP_LOGE(TAG, "DNS bind failed: %d", errno);
        close(sock);
        vTaskDelete(nullptr);
        return;
    }

    ESP_LOGI(TAG, "Captive DNS started");

    uint8_t rx_buf[256];
    while (1) {
        struct sockaddr_in source_addr;
        socklen_t socklen = sizeof(source_addr);
        int len = recvfrom(sock, rx_buf, sizeof(rx_buf), 0, (struct sockaddr *)&source_addr, &socklen);
        if (len < 0) {
            /* Socket closed or error */
            break;
        }
        if (len < sizeof(dns_header_t)) {
            continue;
        }

        dns_header_t *hdr = (dns_header_t *)rx_buf;
        /* Build response in-place after question section */
        uint8_t *ptr = rx_buf + len; // end of packet -> start answer

        /* Question starts after header -> skip QName */
        uint8_t *qname = ((uint8_t *)hdr) + sizeof(dns_header_t);
        uint8_t *q = qname;
        // skip domain labels
        while (q < rx_buf + len && *q) {
            q += (*q) + 1;
        }
        if (q >= rx_buf + len) {
            continue;
        }
        q += 5; // skip null label, qtype, qclass (total 1 + 2 + 2)

        // Prepare answer using pointer to QName (compression pointer 0xC0 0x0C)
        ptr[0] = 0xC0; ptr[1] = 0x0C; // Name pointer
        ptr[2] = 0x00; ptr[3] = 0x01; // Type A
        ptr[4] = 0x00; ptr[5] = 0x01; // Class IN
        ptr[6] = 0x00; ptr[7] = 0x00; ptr[8] = 0x00; ptr[9] = 0x1E; // TTL 30s
        ptr[10] = 0x00; ptr[11] = 0x04; // Data length
        ptr[12] = 192; ptr[13] = 168; ptr[14] = 4; ptr[15] = 1; // 192.168.4.1
        size_t ans_len = 16;

        hdr->flags  = htons(0x8180); // Standard query response, No error
        hdr->ancount = htons(1);

        int tx_len = len + ans_len;
        sendto(sock, rx_buf, tx_len, 0, (struct sockaddr *)&source_addr, sizeof(source_addr));
    }

    /* Cleanup */
    lwip_close(sock);
    ApProvision::_dns_sock = -1;
    vTaskDelete(nullptr);
}

esp_err_t ApProvision::init_dns_server()
{
    if (_dns_task_handle) {
        return ESP_OK;
    }
    BaseType_t ok = xTaskCreate(captive_dns_task, "dns_srv", 4096, nullptr, 3, &_dns_task_handle);
    if (ok != pdPASS) {
        ESP_LOGE(TAG, "Failed to start DNS task");
        return ESP_FAIL;
    }
    return ESP_OK;
}

void ApProvision::deinit_dns_server()
{
    if (_dns_sock >= 0) {
        lwip_close(_dns_sock);
        ApProvision::_dns_sock = -1;
    }
    /* Task will self-delete after socket closure */
    _dns_task_handle = nullptr;
}

/********************** Event handler ***************************/

void ApProvision::wifi_event_handler(void *arg, esp_event_base_t base, int32_t id, void *data)
{
    if (id == WIFI_EVENT_STA_DISCONNECTED) {
        if (_connect_status == ConnectStatus::CONNECTING) {
            wifi_event_sta_disconnected_t *disc = static_cast<wifi_event_sta_disconnected_t *>(data);
            _connect_error_msg = "reason=" + std::to_string(disc ? disc->reason : 0);
            _connect_status = ConnectStatus::FAILED;
            ESP_LOGI(TAG, "Connection failed, reason code: %d", disc ? disc->reason : -1);
        }
    } else if (id == WIFI_EVENT_STA_CONNECTED) {
        ESP_LOGI(TAG, "Connection succeeded");
    }
}

void ApProvision::ip_event_handler(void *arg, esp_event_base_t base, int32_t id, void *data)
{
    if (id != IP_EVENT_STA_GOT_IP) {
        return;
    }

    wifi_ap_record_t ap_info = {};
    if (esp_wifi_sta_get_ap_info(&ap_info) != ESP_OK) {
        return;
    }
    std::string connected_ssid = reinterpret_cast<char *>(ap_info.ssid);
    ESP_LOGI(TAG, "Got IP, connected to: %s", connected_ssid.c_str());

    if (!_target_ssid.empty() && connected_ssid == _target_ssid) {
        ESP_LOGI(TAG, "Provisioning succeeded for SSID %s", _target_ssid.c_str());
        _connect_status = ConnectStatus::SUCCESS;
        if (_cb) {
            _cb(_target_ssid, _target_password);
        }
    }
}

} // namespace esp_brookesia::speaker_apps
